import torch
import torch.nn as nn
from .PConv2d import PConv2d

def create_Pconv_U_start(in_channel, out_channel, kernel_size, stride=1, padding=1,layer=2):
    if layer == 1:
        DPconv = nn.Sequential(
            PConv2d(in_channels=in_channel,out_channels=out_channel,kernel_size=kernel_size,stride=stride,padding=padding),
            nn.ReLU()
        )
    else:
        DPconv = nn.Sequential(
            PConv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=kernel_size, stride=stride,
                    padding=padding),
            nn.BatchNorm2d(num_features=out_channel, momentum=0.5),
            nn.ReLU()
        )
    return DPconv

def concat_(x,y):
    data = torch.cat((x,y),1)
    return data


def create_Pconv_U_end(in_channel, out_channel, kernel_size, stride=1, padding=1):
    UPconv = nn.Sequential(
        PConv2d(in_channels=in_channel,out_channels=out_channel,kernel_size=kernel_size,stride=stride,padding=padding),
        nn.BatchNorm2d(num_features=out_channel),
        nn.LeakyReLU(negative_slope=0.2)
    )
    return UPconv

def create_Upsample(scale_factor):
    return nn.Upsample(scale_factor=scale_factor,mode='nearest')

def create_mask_Conv(kernel_size, stride=1, padding=1):
    mask_conv = nn.Conv2d(in_channels=1,out_channels=1,kernel_size=kernel_size,stride=stride,padding=padding,bias=False)
    mask_conv.weight = nn.Parameter(torch.ones(mask_conv.weight.size()))
    mask_conv.required_grad = False
    return mask_conv


class Pconv_encoder(nn.Module):
    def __init__(self):
        super(Pconv_encoder, self).__init__()

        self.encoder = []
        self.mask_encoder = []
        self.feature = []

        for i in range(8):
            if i == 1:
                self.encoder.append(
                    create_Pconv_U_start(in_channel=3,out_channel=64,kernel_size=7,stride=2,padding=3,layer=i))
                self.mask_encoder.append(create_mask_Conv(kernel_size=7,stride=2,padding=3))
            elif i == 2:
                self.encoder.append(
                    create_Pconv_U_start(in_channel=64,out_channel=128,kernel_size=5,stride=2,padding=2,layer=i))
                self.mask_encoder.append(create_mask_Conv(kernel_size=5,stride=2,padding=2))
            elif i == 3:
                self.encoder.append(
                    create_Pconv_U_start(in_channel=128, out_channel=256, kernel_size=5, stride=2, padding=2, layer=i))
                self.mask_encoder.append(create_mask_Conv(kernel_size=5, stride=2, padding=2))
            elif i == 4:
                self.encoder.append(
                    create_Pconv_U_start(in_channel=256, out_channel=512, kernel_size=3, stride=2, padding=1, layer=i))
                self.mask_encoder.append(create_mask_Conv(kernel_size=3, stride=2, padding=1))
            else:
                self.encoder.append(
                    create_Pconv_U_start(in_channel=512, out_channel=512, kernel_size=3, stride=2, padding=1, layer=i))
                self.mask_encoder.append(create_mask_Conv(kernel_size=3, stride=2, padding=1))

    def forward(self, input, mask):
        self.feature.append(input)
        for i in range(len(self.encoder)):
            input = self.encoder[i](input, mask)
            self.feature.append(input)
            mask = self.mask_encoder[i](mask)
        return self.feature, mask


class Pconv_decoder(nn.Module):
    def __init__(self):
        super(Pconv_decoder, self).__init__()
        self.upsample = create_Upsample(scale_factor=2)
        self.mask_decoder = create_mask_Conv(kernel_size=3,stride=1,padding=1)
        self.decoder = []

        for i in range(8):
            if i in range(0,4):
                self.decoder.append(create_Pconv_U_end(in_channel=1024,out_channel=512,kernel_size=3,stride=1,padding=1))
            elif i == 4:
                self.decoder.append(create_Pconv_U_end(in_channel=768,out_channel=256,kernel_size=3,stride=1,padding=1))
            elif i == 5:
                self.decoder.append(create_Pconv_U_end(in_channel=384, out_channel=128, kernel_size=3, stride=1,padding=1))
            elif i == 6:
                self.decoder.append(create_Pconv_U_end(in_channel=192, out_channel=64, kernel_size=3, stride=1,padding=1))
            elif i == 7:
                self.decoder.append(PConv2d(in_channels=67, out_channels=3, kernel_size=3, stride=1,padding=1))


    def forward(self, input, mask):
        input.reverse()
        x = input[0]
        for i in range(1, len(input)):
            x = self.upsample(x)
            mask = self.upsample(mask)
            x = concat_(x,input[i])
            x = self.decoder[i-1](x,mask)
            mask = self.mask_decoder(mask)
        return x, mask




